Module Module1

    ' PC crashed, leaving the removable device files missing. CHKDSK /F recovers the lost file chains as FILE????.CHK files.
    ' scan the file header to determine the filetype, and rename the file accordingly.

    Private buffersize As Integer = 100
    Dim keys As New List(Of String)
    Dim sigs As New List(Of String)


    'Const logoTop As String = " (                   (                                        " & ControlChars.CrLf & _
    '                    " )\ )      (         )\ )                                     " & ControlChars.CrLf & _
    '                    "(()/(  (   )\   (   (()/(   (            )      (   (    (    " & ControlChars.CrLf & _
    '                    " /(_)) )\ ((_) ))\   /(_)) ))\  (   (   /((    ))\  )(   )\ ) " & ControlChars.CrLf & _
    '                    "(_))_|((_) _  /((_) (_))  /((_) )\  )\ (_))\  /((_)(()\ (()/( " & ControlChars.CrLf & _
    '                    "| |_   (_)| |(_))   | _ \(_))  ((_)((_)_)((_)(_))   ((_) )(_))" & ControlChars.CrLf & _
    '                    "| __|  | || |/ -_)  |   // -_)/ _|/ _ \\ V / / -_) | '_|| || |" & ControlChars.CrLf & _
    '                    "|_|    |_||_|\___|  |_|_\\___|\__|\___/ \_/  \___| |_|   \_, |" & ControlChars.CrLf & _
    '                    "                                                         |__/ " & ControlChars.CrLf & _
    '                    " o---------------------------------------------------------o" & ControlChars.CrLf & _
    '                    " | Did your pc crashed? Did you have a removable drive     |" & ControlChars.CrLf & _
    '                    " | connected, and did your files go 'missing'?             |" & ControlChars.CrLf & _
    '                    " | You ran CHKDSK /F and your files recovered as.. *gulp*  |" & ControlChars.CrLf & _
    '                    " |                                                         |" & ControlChars.CrLf & _
    '                    " |         FILE0000.CHK up to FILE9999.CHK ?               |" & ControlChars.CrLf & _
    '                    " |                                                         |" & ControlChars.CrLf & _
    '                    " | This tool will scan the file header, determine the      |" & ControlChars.CrLf & _
    '                    " | file type and rename the file extension, bringing       |" & ControlChars.CrLf & _
    '                    " | you one step closer to recovering the files you need.   |" & ControlChars.CrLf & _
    '                    " |                                                         |" & ControlChars.CrLf & _
    '                    " o---------------------------------------------------------o"

    Const logoTop As String = " (                   (                                        " & ControlChars.CrLf & _
                            " )\ )      (         )\ )                                     " & ControlChars.CrLf & _
                            "(()/(  (   )\   (   (()/(   (            )      (   (    (    " & ControlChars.CrLf & _
                            " /(_)) )\ ((_) ))\   /(_)) ))\  (   (   /((    ))\  )(   )\ ) " & ControlChars.CrLf & _
                            "(_))_|((_) _  /((_) (_))  /((_) )\  )\ (_))\  /((_)(()\ (()/( "


    Const logoMiddle As String = "| |_   (_)| |(_))   | _ \(_))  ((_)((_)_)((_)(_))   ((_) )(_))" & ControlChars.CrLf & _
                                "| __|  | || |/ -_)  |   // -_)/ _|/ _ \\ V / / -_) | '_|| || |" & ControlChars.CrLf & _
                                "|_|    |_||_|\___|  |_|_\\___|\__|\___/ \_/  \___| |_|   \_, |" & ControlChars.CrLf & _
                                "                                                         |__/ "


    Const logoBottom As String = " o---------------------------------------------------------o" & ControlChars.CrLf & _
                                " | Did your pc crash? Did you have a removable drive       |" & ControlChars.CrLf & _
                                " | connected, and did your files go 'missing'?             |" & ControlChars.CrLf & _
                                " | You ran CHKDSK /F and your files recovered as.. *gulp*  |" & ControlChars.CrLf & _
                                " |                                                         |" & ControlChars.CrLf & _
                                " |         FILE0000.CHK up to FILE9999.CHK ?               |" & ControlChars.CrLf & _
                                " |                                                         |" & ControlChars.CrLf & _
                                " | This tool will scan the file header, determine the      |" & ControlChars.CrLf & _
                                " | file type and rename the file extension, bringing       |" & ControlChars.CrLf & _
                                " | you one step closer to recovering the files you need.   |" & ControlChars.CrLf & _
                                " |                                                         |" & ControlChars.CrLf & _
                                " o---------------------------------------------------------o"



    ''' <summary>
    ''' create mapping file
    ''' </summary>
    ''' <returns></returns>
    ''' <remarks></remarks>
    Private Function createFile() As Boolean
        Try

            Dim f As String = IO.Path.Combine(My.Application.Info.DirectoryPath, "headers.txt")

            If Not IO.File.Exists(f) Then
                ' create the file if it doesnt exist
                Dim writer As New IO.StreamWriter(f, False)
                writer.WriteLine("########################################################")
                writer.WriteLine("# Header definition file")
                writer.WriteLine("# Each line contains a mapping, in the format <ext>=<sig>")
                writer.WriteLine("# <ext> = file extension to map")
                writer.WriteLine("# <sig> = header signature to look for in the first 100 bytes")
                writer.WriteLine("# ")
                writer.WriteLine("########################################################")
                writer.WriteLine("# Buffersize determines how many bytes of the file header to scan for signatures")
                writer.WriteLine("buffersize={0}", buffersize)
                writer.WriteLine("# ")
                writer.WriteLine("########################################################")
                writer.WriteLine("# Note: I will only process files with the .CHK extension, the recovered file extension.")
                writer.WriteLine("########################################################")
                writer.WriteLine("exe=MZ")
                writer.WriteLine("pdf=pdf")
                writer.WriteLine("wav=RIFF")
                writer.WriteLine("mp3=ID3")
                writer.WriteLine("zip=PK")
                writer.WriteLine("jpg=JFIF")
                writer.WriteLine("html=HTML")
                writer.Flush()
                writer.Close()
            End If


            ' read the file
            Dim reader As New IO.StreamReader(f)
            Dim line As String = reader.ReadLine

            While Not IsNothing(line)

                ' set the buffer
                If line.StartsWith("buffersize") Then

                    ' extract the buffer size value
                    Dim index As Integer = line.IndexOf("="c)

                    If (index > -1) Then
                        Dim size As Object = line.Substring(index + 1)
                        If IsNumeric(size) Then
                            buffersize = CInt(size)
                        Else
                            buffersize = -1
                        End If
                    End If

                Else

                    ' if not a comment line
                    If Not line.StartsWith("#") Then

                        ' lookup the map
                        Dim index As Integer = line.IndexOf("="c)

                        If (index > -1) Then

                            keys.Add(line.Substring(0, index))
                            sigs.Add(line.Substring(index + 1))

                        End If

                    End If

                End If

                ' read the next line
                line = reader.ReadLine

            End While

            Return True

        Catch ex As Exception
            Console.WriteLine("Error trying to create the default header file: {0}", ex.Message)
            Console.WriteLine("press <enter> to exit")
            Console.ReadLine()
        End Try
    End Function

    Sub Main()

        If Not createFile() Then Exit Sub

        Dim ipod As String

        Console.ForegroundColor = ConsoleColor.DarkYellow
        Console.WriteLine(logoTop)
        Console.ForegroundColor = ConsoleColor.Red
        Console.WriteLine(logoMiddle)
        Console.ForegroundColor = ConsoleColor.White
        Console.WriteLine(logoBottom)

        Console.ForegroundColor = ConsoleColor.Gray
        Console.WriteLine()
        Console.WriteLine("[Read {0} header mappings]", keys.Count)

        ' check for a invalid buffersize
        If (buffersize = -1) Then
            Console.WriteLine("[The specified buffer size is invalid, please supply a valid number]")
            Console.WriteLine("press <enter> to exit")
            Console.ReadLine()
            Exit Sub
        Else
            Console.WriteLine("[Buffer size set to {0} bytes]", buffersize)
        End If

        Console.WriteLine("(Hint: delete the headers file to have it reset to defaults)")
        Console.WriteLine()



        ' start the UI area
        Console.ForegroundColor = ConsoleColor.White

        ' read the path to process
        Console.Write("Enter the path to the recovered files before we begin: ")
        ipod = Console.ReadLine

        If Not IO.Directory.Exists(ipod) Then
            If (ipod.Length > 0) Then
                ' show a message if it looks like the user tried to enter a path
                Console.WriteLine("The path you entered does not exist.")
                Console.WriteLine("press <enter> to exit")
                Console.ReadLine()
            End If
            Exit Sub
        End If

        Dim filemask As String = "*.chk"
        Dim total As Integer = 0
        Dim fixed As Integer = 0

        Dim files() As String = IO.Directory.GetFiles(ipod, filemask)

        total = files.Length

        Console.WriteLine("found {0} files", total)

        For Each filename As String In files

            If determindeFileType(filename) Then fixed += 1

        Next

        Console.WriteLine("Files Fixed: {0}", fixed)
        Console.WriteLine("press <enter> to exit")
        Console.ReadLine()

    End Sub


    Private Function determindeFileType(ByVal f As String) As Boolean

        Try

            Dim newExtension As String = Nothing
            Dim buff(buffersize - 1) As Byte

            ' open the filestream and read the header into a buffer
            Dim sin As New IO.FileStream(f, IO.FileMode.Open)

            Dim bytesRead As Integer = 0

            bytesRead = sin.Read(buff, 0, buffersize)

            ' close the stream
            sin.Close()

            ' convert the buffer to a string
            Dim bufferstring As String = "" 'System.Text.ASCIIEncoding.ASCII.GetString(buff)

            ' clean up the buffer
            For index As Integer = 0 To bytesRead - 1
                Dim c As Char = ChrW(buff(index))
                If Char.IsLetterOrDigit(c) Or Char.IsPunctuation(c) Or Char.IsSymbol(c) Or Char.IsWhiteSpace(c) Then
                    bufferstring += ChrW(buff(index))
                End If
            Next

            ' loop through the list of header mappings and look for a matching pattern
            For Each sig As String In sigs
                If bufferstring.Contains(sig) Then
                    newExtension = keys(sigs.IndexOf(sig))
                    Exit For
                End If
            Next

            ' rename the file to the new extension
            If Not IsNothing(newExtension) Then

                ' ensure the extension format is correct
                If Not newExtension.StartsWith(".") Then newExtension = "." & newExtension

                ' rename the file using the fileinfo object
                Dim fi As New IO.FileInfo(f)
                Microsoft.VisualBasic.FileSystem.Rename(f, IO.Path.ChangeExtension(f, newExtension))

                ' we found the extension! success!
                Return True

            End If

        Catch ex As Exception

        End Try

    End Function

End Module
